import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import tensorflow as tf
from sklearn.decomposition import PCA
import os
import timeit

PATCH_SIZE = 32
#Size of random sample from entire input space
BATCH_SIZE = 1944
#Use choices 'modindex_nat', 'modindex_noise', 'fg/', 'tex/' and 'line/' depending on the choice of input below
OUT_FOLDER = 'fg/'

start = timeit.default_timer()

if not os.path.exists(OUT_FOLDER):
    os.makedirs(OUT_FOLDER)

filters = np.load('gabors.npy').astype(np.float32)

#Choose input stimuli for either modindex calculation or classification

#Mod index: naturalistic textures
#samples = np.load('modindex_texturesfamily/nattextures32.npy')

#Mod index: noise textures
#samples = np.load('modindex_texturesfamily/noisetextures32.npy')

#Figure-ground
samples = np.load('fg32.npy')

#Texture
#samples = np.load('textures32.npy')

#Line angle
#samples = np.load('lineangle32.npy')

samples = samples.reshape((samples.shape[0], PATCH_SIZE, PATCH_SIZE, 1))

patchCount = samples.shape[0]

v1Simple = np.zeros((patchCount, 6, 6, 3, 12, 2))
v1Complex = np.zeros((patchCount, 1296))
anglesFile = np.zeros((patchCount, 6, 6, 3, 12, 1))

loopCount = int(np.ceil(patchCount / BATCH_SIZE))

filters = np.reshape(filters, (12, 12, 1, -1))

with tf.device('/gpu:0'):
	for i in range(loopCount):

		bStart = i * BATCH_SIZE
		bEnd = bStart + BATCH_SIZE

		batch = samples[bStart:bEnd, :, :, :]

		filtersTensor = tf.convert_to_tensor(filters)

		v1Responses = tf.nn.conv2d(batch, filtersTensor, [1, 4, 4, 1], 'VALID')

		v1Responses = tf.reshape(v1Responses, [-1, 6, 6, 3, 12, 2])

		pair0, pair1 = tf.split(v1Responses, num_or_size_splits = 2, axis = -1)

		angles = tf.atan2(pair1, pair0)

		v1CResponses = tf.sqrt(tf.reduce_sum(tf.square(v1Responses), axis = -1))

		v1CResponses = tf.reshape(v1CResponses, [v1CResponses.shape[0], -1])

		v1Simple[bStart:bEnd, :] = v1Responses[:]
		v1Complex[bStart:bEnd, :] = v1CResponses[:]
		anglesFile[bStart:bEnd, :, :, :, :, :] = angles[:]

v1ComplexMean = np.mean(v1Complex, axis = 1, keepdims = True)

v1Complex -= v1ComplexMean

forwardPCA = np.load('imnet/forwardPCA.npy')

pcaTransformed = np.dot(v1Complex, forwardPCA.T)

#Choose output directory for stimuli choice

#Figure-ground
np.save(OUT_FOLDER + 'v1Simple.npy', v1Simple)
np.save(OUT_FOLDER + 'v1Complex.npy', v1Complex)
np.save(OUT_FOLDER + 'angles.npy', anglesFile)
np.save(OUT_FOLDER + 'pcaTransformed.npy', pcaTransformed)
np.save(OUT_FOLDER + 'v1cmean.npy', v1ComplexMean)

stop = timeit.default_timer()

print('Time: ', stop - start)
